import asyncio
import os
from fastapi import APIRouter, WebSocket, UploadFile, File, Form, Depends
from fastapi.responses import StreamingResponse
from .stream_config import StreamConfig, MediaType
from .orchestrator import StreamOrchestrator
from utils.logger import setup_logger
logger = setup_logger()


from utils.file_handler import save_temp_file, get_mime_type, create_stream_response
router = APIRouter()


@router.post("/process")
async def process_stream(config: StreamConfig):
    orchestrator = StreamOrchestrator(config)
    async def input_data_generator():
        for i in range(10):
            yield f"Input data {i}"
            await asyncio.sleep(0.1)

    # 根据配置的媒体类型设置响应头
    media_type = config.media_type.value if config.media_type else "application/octet-stream"
    
    # 根据输出类型调整媒体类型
    if config.output_type.value.endswith('_stream'):
        if config.output_type.value == 'text_stream':
            media_type = MediaType.TEXT.value
        elif config.output_type.value == 'audio_stream':
            media_type = MediaType.AUDIO.value
        elif config.output_type.value == 'video_stream':
            media_type = MediaType.VIDEO.value
    
    return StreamingResponse(
        orchestrator.process(input_data_generator()),
        media_type=media_type
    )

@router.post("/process/file")
async def process_file_stream(
    file: UploadFile = File(...),
    config_data: str = Form(...)
):
    import json
    config_dict = json.loads(config_data)
    config = StreamConfig(**config_dict)
    
    # 保存上传的文件
    temp_file_path = await save_temp_file(file)
    
    # 创建编排器
    orchestrator = StreamOrchestrator(config)
    
    # 确定媒体类型
    media_type = get_mime_type(temp_file_path)
    if config.media_type:
        media_type = config.media_type.value
    
    # 根据输出类型调整媒体类型
    if config.output_type.value.endswith('_stream'):
        if config.output_type.value == 'text_stream':
            media_type = MediaType.TEXT.value
        elif config.output_type.value == 'audio_stream':
            media_type = MediaType.AUDIO.value
        elif config.output_type.value == 'video_stream':
            media_type = MediaType.VIDEO.value
    
    # 处理流
    return StreamingResponse(
        orchestrator.process(temp_file_path),
        media_type=media_type
    )

@router.websocket("/ws/process")
async def websocket_process(websocket: WebSocket):
    await websocket.accept()
    try:
        config_data = await websocket.receive_json()
        config = StreamConfig(**config_data)
        orchestrator = StreamOrchestrator(config)

        async def input_data_generator():
            while True:
                data = await websocket.receive_text()
                yield data

        async for output in orchestrator.process(input_data_generator()):
            await websocket.send_bytes(output)
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
        await websocket.close()
    finally:
        await websocket.close()